{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\nFinetuning Torchvision Models\n=============================\n\n**Author:** `Nathan Inkawhich <https://github.com/inkawhich>`__\n\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In this tutorial we will take a deeper look at how to finetune and\nfeature extract the `torchvision\nmodels <https://pytorch.org/docs/stable/torchvision/models.html>`__, all\nof which have been pretrained on the 1000-class Imagenet dataset. This\ntutorial will give an indepth look at how to work with several modern\nCNN architectures, and will build an intuition for finetuning any\nPyTorch model. Since each model architecture is different, there is no\nboilerplate finetuning code that will work in all scenarios. Rather, the\nresearcher must look at the existing architecture and make custom\nadjustments for each model.\n\nIn this document we will perform two types of transfer learning:\nfinetuning and feature extraction. In **finetuning**, we start with a\npretrained model and update *all* of the model\u2019s parameters for our new\ntask, in essence retraining the whole model. In **feature extraction**,\nwe start with a pretrained model and only update the final layer weights\nfrom which we derive predictions. It is called feature extraction\nbecause we use the pretrained CNN as a fixed feature-extractor, and only\nchange the output layer. For more technical information about transfer\nlearning see `here <https://cs231n.github.io/transfer-learning/>`__ and\n`here <https://ruder.io/transfer-learning/>`__.\n\nIn general both transfer learning methods follow the same few steps:\n\n-  Initialize the pretrained model\n-  Reshape the final layer(s) to have the same number of outputs as the\n   number of classes in the new dataset\n-  Define for the optimization algorithm which parameters we want to\n   update during training\n-  Run the training step\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from __future__ import print_function \nfrom __future__ import division\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport numpy as np\nimport torchvision\nfrom torchvision import datasets, models, transforms\nimport matplotlib.pyplot as plt\nimport time\nimport os\nimport copy\nprint(\"PyTorch Version: \",torch.__version__)\nprint(\"Torchvision Version: \",torchvision.__version__)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Inputs\n------\n\nHere are all of the parameters to change for the run. We will use the\n*hymenoptera_data* dataset which can be downloaded\n`here <https://download.pytorch.org/tutorial/hymenoptera_data.zip>`__.\nThis dataset contains two classes, **bees** and **ants**, and is\nstructured such that we can use the\n`ImageFolder <https://pytorch.org/docs/stable/torchvision/datasets.html#torchvision.datasets.ImageFolder>`__\ndataset, rather than writing our own custom dataset. Download the data\nand set the ``data_dir`` input to the root directory of the dataset. The\n``model_name`` input is the name of the model you wish to use and must\nbe selected from this list:\n\n::\n\n   [resnet, alexnet, vgg, squeezenet, densenet, inception]\n\nThe other inputs are as follows: ``num_classes`` is the number of\nclasses in the dataset, ``batch_size`` is the batch size used for\ntraining and may be adjusted according to the capability of your\nmachine, ``num_epochs`` is the number of training epochs we want to run,\nand ``feature_extract`` is a boolean that defines if we are finetuning\nor feature extracting. If ``feature_extract = False``, the model is\nfinetuned and all model parameters are updated. If\n``feature_extract = True``, only the last layer parameters are updated,\nthe others remain fixed.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Top level data directory. Here we assume the format of the directory conforms \n#   to the ImageFolder structure\ndata_dir = \"./data/hymenoptera_data\"\n\n# Models to choose from [resnet, alexnet, vgg, squeezenet, densenet, inception]\nmodel_name = \"squeezenet\"\n\n# Number of classes in the dataset\nnum_classes = 2\n\n# Batch size for training (change depending on how much memory you have)\nbatch_size = 8\n\n# Number of epochs to train for \nnum_epochs = 15\n\n# Flag for feature extracting. When False, we finetune the whole model, \n#   when True we only update the reshaped layer params\nfeature_extract = True"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Helper Functions\n----------------\n\nBefore we write the code for adjusting the models, lets define a few\nhelper functions.\n\nModel Training and Validation Code\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nThe ``train_model`` function handles the training and validation of a\ngiven model. As input, it takes a PyTorch model, a dictionary of\ndataloaders, a loss function, an optimizer, a specified number of epochs\nto train and validate for, and a boolean flag for when the model is an\nInception model. The *is_inception* flag is used to accomodate the\n*Inception v3* model, as that architecture uses an auxiliary output and\nthe overall model loss respects both the auxiliary output and the final\noutput, as described\n`here <https://discuss.pytorch.org/t/how-to-optimize-inception-model-with-auxiliary-classifiers/7958>`__.\nThe function trains for the specified number of epochs and after each\nepoch runs a full validation step. It also keeps track of the best\nperforming model (in terms of validation accuracy), and at the end of\ntraining returns the best performing model. After each epoch, the\ntraining and validation accuracies are printed.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def train_model(model, dataloaders, criterion, optimizer, num_epochs=25, is_inception=False):\n    since = time.time()\n\n    val_acc_history = []\n    \n    best_model_wts = copy.deepcopy(model.state_dict())\n    best_acc = 0.0\n\n    for epoch in range(num_epochs):\n        print('Epoch {}/{}'.format(epoch, num_epochs - 1))\n        print('-' * 10)\n\n        # Each epoch has a training and validation phase\n        for phase in ['train', 'val']:\n            if phase == 'train':\n                model.train()  # Set model to training mode\n            else:\n                model.eval()   # Set model to evaluate mode\n\n            running_loss = 0.0\n            running_corrects = 0\n\n            # Iterate over data.\n            for inputs, labels in dataloaders[phase]:\n                inputs = inputs.to(device)\n                labels = labels.to(device)\n\n                # zero the parameter gradients\n                optimizer.zero_grad()\n\n                # forward\n                # track history if only in train\n                with torch.set_grad_enabled(phase == 'train'):\n                    # Get model outputs and calculate loss\n                    # Special case for inception because in training it has an auxiliary output. In train\n                    #   mode we calculate the loss by summing the final output and the auxiliary output\n                    #   but in testing we only consider the final output.\n                    if is_inception and phase == 'train':\n                        # From https://discuss.pytorch.org/t/how-to-optimize-inception-model-with-auxiliary-classifiers/7958\n                        outputs, aux_outputs = model(inputs)\n                        loss1 = criterion(outputs, labels)\n                        loss2 = criterion(aux_outputs, labels)\n                        loss = loss1 + 0.4*loss2\n                    else:\n                        outputs = model(inputs)\n                        loss = criterion(outputs, labels)\n\n                    _, preds = torch.max(outputs, 1)\n\n                    # backward + optimize only if in training phase\n                    if phase == 'train':\n                        loss.backward()\n                        optimizer.step()\n\n                # statistics\n                running_loss += loss.item() * inputs.size(0)\n                running_corrects += torch.sum(preds == labels.data)\n\n            epoch_loss = running_loss / len(dataloaders[phase].dataset)\n            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)\n\n            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))\n\n            # deep copy the model\n            if phase == 'val' and epoch_acc > best_acc:\n                best_acc = epoch_acc\n                best_model_wts = copy.deepcopy(model.state_dict())\n            if phase == 'val':\n                val_acc_history.append(epoch_acc)\n\n        print()\n\n    time_elapsed = time.time() - since\n    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))\n    print('Best val Acc: {:4f}'.format(best_acc))\n\n    # load best model weights\n    model.load_state_dict(best_model_wts)\n    return model, val_acc_history"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Set Model Parameters\u2019 .requires_grad attribute\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nThis helper function sets the ``.requires_grad`` attribute of the\nparameters in the model to False when we are feature extracting. By\ndefault, when we load a pretrained model all of the parameters have\n``.requires_grad=True``, which is fine if we are training from scratch\nor finetuning. However, if we are feature extracting and only want to\ncompute gradients for the newly initialized layer then we want all of\nthe other parameters to not require gradients. This will make more sense\nlater.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def set_parameter_requires_grad(model, feature_extracting):\n    if feature_extracting:\n        for param in model.parameters():\n            param.requires_grad = False"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Initialize and Reshape the Networks\n-----------------------------------\n\nNow to the most interesting part. Here is where we handle the reshaping\nof each network. Note, this is not an automatic procedure and is unique\nto each model. Recall, the final layer of a CNN model, which is often\ntimes an FC layer, has the same number of nodes as the number of output\nclasses in the dataset. Since all of the models have been pretrained on\nImagenet, they all have output layers of size 1000, one node for each\nclass. The goal here is to reshape the last layer to have the same\nnumber of inputs as before, AND to have the same number of outputs as\nthe number of classes in the dataset. In the following sections we will\ndiscuss how to alter the architecture of each model individually. But\nfirst, there is one important detail regarding the difference between\nfinetuning and feature-extraction.\n\nWhen feature extracting, we only want to update the parameters of the\nlast layer, or in other words, we only want to update the parameters for\nthe layer(s) we are reshaping. Therefore, we do not need to compute the\ngradients of the parameters that we are not changing, so for efficiency\nwe set the .requires_grad attribute to False. This is important because\nby default, this attribute is set to True. Then, when we initialize the\nnew layer and by default the new parameters have ``.requires_grad=True``\nso only the new layer\u2019s parameters will be updated. When we are\nfinetuning we can leave all of the .required_grad\u2019s set to the default\nof True.\n\nFinally, notice that inception_v3 requires the input size to be\n(299,299), whereas all of the other models expect (224,224).\n\nResnet\n~~~~~~\n\nResnet was introduced in the paper `Deep Residual Learning for Image\nRecognition <https://arxiv.org/abs/1512.03385>`__. There are several\nvariants of different sizes, including Resnet18, Resnet34, Resnet50,\nResnet101, and Resnet152, all of which are available from torchvision\nmodels. Here we use Resnet18, as our dataset is small and only has two\nclasses. When we print the model, we see that the last layer is a fully\nconnected layer as shown below:\n\n::\n\n   (fc): Linear(in_features=512, out_features=1000, bias=True) \n\nThus, we must reinitialize ``model.fc`` to be a Linear layer with 512\ninput features and 2 output features with:\n\n::\n\n   model.fc = nn.Linear(512, num_classes)\n\nAlexnet\n~~~~~~~\n\nAlexnet was introduced in the paper `ImageNet Classification with Deep\nConvolutional Neural\nNetworks <https://papers.nips.cc/paper/4824-imagenet-classification-with-deep-convolutional-neural-networks.pdf>`__\nand was the first very successful CNN on the ImageNet dataset. When we\nprint the model architecture, we see the model output comes from the 6th\nlayer of the classifier\n\n::\n\n   (classifier): Sequential(\n       ...\n       (6): Linear(in_features=4096, out_features=1000, bias=True)\n    ) \n\nTo use the model with our dataset we reinitialize this layer as\n\n::\n\n   model.classifier[6] = nn.Linear(4096,num_classes)\n\nVGG\n~~~\n\nVGG was introduced in the paper `Very Deep Convolutional Networks for\nLarge-Scale Image Recognition <https://arxiv.org/pdf/1409.1556.pdf>`__.\nTorchvision offers eight versions of VGG with various lengths and some\nthat have batch normalizations layers. Here we use VGG-11 with batch\nnormalization. The output layer is similar to Alexnet, i.e.\n\n::\n\n   (classifier): Sequential(\n       ...\n       (6): Linear(in_features=4096, out_features=1000, bias=True)\n    )\n\nTherefore, we use the same technique to modify the output layer\n\n::\n\n   model.classifier[6] = nn.Linear(4096,num_classes)\n\nSqueezenet\n~~~~~~~~~~\n\nThe Squeeznet architecture is described in the paper `SqueezeNet:\nAlexNet-level accuracy with 50x fewer parameters and <0.5MB model\nsize <https://arxiv.org/abs/1602.07360>`__ and uses a different output\nstructure than any of the other models shown here. Torchvision has two\nversions of Squeezenet, we use version 1.0. The output comes from a 1x1\nconvolutional layer which is the 1st layer of the classifier:\n\n::\n\n   (classifier): Sequential(\n       (0): Dropout(p=0.5)\n       (1): Conv2d(512, 1000, kernel_size=(1, 1), stride=(1, 1))\n       (2): ReLU(inplace)\n       (3): AvgPool2d(kernel_size=13, stride=1, padding=0)\n    ) \n\nTo modify the network, we reinitialize the Conv2d layer to have an\noutput feature map of depth 2 as\n\n::\n\n   model.classifier[1] = nn.Conv2d(512, num_classes, kernel_size=(1,1), stride=(1,1))\n\nDensenet\n~~~~~~~~\n\nDensenet was introduced in the paper `Densely Connected Convolutional\nNetworks <https://arxiv.org/abs/1608.06993>`__. Torchvision has four\nvariants of Densenet but here we only use Densenet-121. The output layer\nis a linear layer with 1024 input features:\n\n::\n\n   (classifier): Linear(in_features=1024, out_features=1000, bias=True) \n\nTo reshape the network, we reinitialize the classifier\u2019s linear layer as\n\n::\n\n   model.classifier = nn.Linear(1024, num_classes)\n\nInception v3\n~~~~~~~~~~~~\n\nFinally, Inception v3 was first described in `Rethinking the Inception\nArchitecture for Computer\nVision <https://arxiv.org/pdf/1512.00567v1.pdf>`__. This network is\nunique because it has two output layers when training. The second output\nis known as an auxiliary output and is contained in the AuxLogits part\nof the network. The primary output is a linear layer at the end of the\nnetwork. Note, when testing we only consider the primary output. The\nauxiliary output and primary output of the loaded model are printed as:\n\n::\n\n   (AuxLogits): InceptionAux(\n       ...\n       (fc): Linear(in_features=768, out_features=1000, bias=True)\n    )\n    ...\n   (fc): Linear(in_features=2048, out_features=1000, bias=True)\n\nTo finetune this model we must reshape both layers. This is accomplished\nwith the following\n\n::\n\n   model.AuxLogits.fc = nn.Linear(768, num_classes)\n   model.fc = nn.Linear(2048, num_classes)\n\nNotice, many of the models have similar output structures, but each must\nbe handled slightly differently. Also, check out the printed model\narchitecture of the reshaped network and make sure the number of output\nfeatures is the same as the number of classes in the dataset.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True):\n    # Initialize these variables which will be set in this if statement. Each of these\n    #   variables is model specific.\n    model_ft = None\n    input_size = 0\n\n    if model_name == \"resnet\":\n        \"\"\" Resnet18\n        \"\"\"\n        model_ft = models.resnet18(pretrained=use_pretrained)\n        set_parameter_requires_grad(model_ft, feature_extract)\n        num_ftrs = model_ft.fc.in_features\n        model_ft.fc = nn.Linear(num_ftrs, num_classes)\n        input_size = 224\n\n    elif model_name == \"alexnet\":\n        \"\"\" Alexnet\n        \"\"\"\n        model_ft = models.alexnet(pretrained=use_pretrained)\n        set_parameter_requires_grad(model_ft, feature_extract)\n        num_ftrs = model_ft.classifier[6].in_features\n        model_ft.classifier[6] = nn.Linear(num_ftrs,num_classes)\n        input_size = 224\n\n    elif model_name == \"vgg\":\n        \"\"\" VGG11_bn\n        \"\"\"\n        model_ft = models.vgg11_bn(pretrained=use_pretrained)\n        set_parameter_requires_grad(model_ft, feature_extract)\n        num_ftrs = model_ft.classifier[6].in_features\n        model_ft.classifier[6] = nn.Linear(num_ftrs,num_classes)\n        input_size = 224\n\n    elif model_name == \"squeezenet\":\n        \"\"\" Squeezenet\n        \"\"\"\n        model_ft = models.squeezenet1_0(pretrained=use_pretrained)\n        set_parameter_requires_grad(model_ft, feature_extract)\n        model_ft.classifier[1] = nn.Conv2d(512, num_classes, kernel_size=(1,1), stride=(1,1))\n        model_ft.num_classes = num_classes\n        input_size = 224\n\n    elif model_name == \"densenet\":\n        \"\"\" Densenet\n        \"\"\"\n        model_ft = models.densenet121(pretrained=use_pretrained)\n        set_parameter_requires_grad(model_ft, feature_extract)\n        num_ftrs = model_ft.classifier.in_features\n        model_ft.classifier = nn.Linear(num_ftrs, num_classes) \n        input_size = 224\n\n    elif model_name == \"inception\":\n        \"\"\" Inception v3 \n        Be careful, expects (299,299) sized images and has auxiliary output\n        \"\"\"\n        model_ft = models.inception_v3(pretrained=use_pretrained)\n        set_parameter_requires_grad(model_ft, feature_extract)\n        # Handle the auxilary net\n        num_ftrs = model_ft.AuxLogits.fc.in_features\n        model_ft.AuxLogits.fc = nn.Linear(num_ftrs, num_classes)\n        # Handle the primary net\n        num_ftrs = model_ft.fc.in_features\n        model_ft.fc = nn.Linear(num_ftrs,num_classes)\n        input_size = 299\n\n    else:\n        print(\"Invalid model name, exiting...\")\n        exit()\n    \n    return model_ft, input_size\n\n# Initialize the model for this run\nmodel_ft, input_size = initialize_model(model_name, num_classes, feature_extract, use_pretrained=True)\n\n# Print the model we just instantiated\nprint(model_ft)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Load Data\n---------\n\nNow that we know what the input size must be, we can initialize the data\ntransforms, image datasets, and the dataloaders. Notice, the models were\npretrained with the hard-coded normalization values, as described\n`here <https://pytorch.org/docs/master/torchvision/models.html>`__.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Data augmentation and normalization for training\n# Just normalization for validation\ndata_transforms = {\n    'train': transforms.Compose([\n        transforms.RandomResizedCrop(input_size),\n        transforms.RandomHorizontalFlip(),\n        transforms.ToTensor(),\n        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n    ]),\n    'val': transforms.Compose([\n        transforms.Resize(input_size),\n        transforms.CenterCrop(input_size),\n        transforms.ToTensor(),\n        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n    ]),\n}\n\nprint(\"Initializing Datasets and Dataloaders...\")\n\n# Create training and validation datasets\nimage_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'val']}\n# Create training and validation dataloaders\ndataloaders_dict = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size, shuffle=True, num_workers=4) for x in ['train', 'val']}\n\n# Detect if we have a GPU available\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Create the Optimizer\n--------------------\n\nNow that the model structure is correct, the final step for finetuning\nand feature extracting is to create an optimizer that only updates the\ndesired parameters. Recall that after loading the pretrained model, but\nbefore reshaping, if ``feature_extract=True`` we manually set all of the\nparameter\u2019s ``.requires_grad`` attributes to False. Then the\nreinitialized layer\u2019s parameters have ``.requires_grad=True`` by\ndefault. So now we know that *all parameters that have\n.requires_grad=True should be optimized.* Next, we make a list of such\nparameters and input this list to the SGD algorithm constructor.\n\nTo verify this, check out the printed parameters to learn. When\nfinetuning, this list should be long and include all of the model\nparameters. However, when feature extracting this list should be short\nand only include the weights and biases of the reshaped layers.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Send the model to GPU\nmodel_ft = model_ft.to(device)\n\n# Gather the parameters to be optimized/updated in this run. If we are\n#  finetuning we will be updating all parameters. However, if we are \n#  doing feature extract method, we will only update the parameters\n#  that we have just initialized, i.e. the parameters with requires_grad\n#  is True.\nparams_to_update = model_ft.parameters()\nprint(\"Params to learn:\")\nif feature_extract:\n    params_to_update = []\n    for name,param in model_ft.named_parameters():\n        if param.requires_grad == True:\n            params_to_update.append(param)\n            print(\"\\t\",name)\nelse:\n    for name,param in model_ft.named_parameters():\n        if param.requires_grad == True:\n            print(\"\\t\",name)\n\n# Observe that all parameters are being optimized\noptimizer_ft = optim.SGD(params_to_update, lr=0.001, momentum=0.9)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Run Training and Validation Step\n--------------------------------\n\nFinally, the last step is to setup the loss for the model, then run the\ntraining and validation function for the set number of epochs. Notice,\ndepending on the number of epochs this step may take a while on a CPU.\nAlso, the default learning rate is not optimal for all of the models, so\nto achieve maximum accuracy it would be necessary to tune for each model\nseparately.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Setup the loss fxn\ncriterion = nn.CrossEntropyLoss()\n\n# Train and evaluate\nmodel_ft, hist = train_model(model_ft, dataloaders_dict, criterion, optimizer_ft, num_epochs=num_epochs, is_inception=(model_name==\"inception\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Comparison with Model Trained from Scratch\n------------------------------------------\n\nJust for fun, lets see how the model learns if we do not use transfer\nlearning. The performance of finetuning vs.\u00a0feature extracting depends\nlargely on the dataset but in general both transfer learning methods\nproduce favorable results in terms of training time and overall accuracy\nversus a model trained from scratch.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Initialize the non-pretrained version of the model used for this run\nscratch_model,_ = initialize_model(model_name, num_classes, feature_extract=False, use_pretrained=False)\nscratch_model = scratch_model.to(device)\nscratch_optimizer = optim.SGD(scratch_model.parameters(), lr=0.001, momentum=0.9)\nscratch_criterion = nn.CrossEntropyLoss()\n_,scratch_hist = train_model(scratch_model, dataloaders_dict, scratch_criterion, scratch_optimizer, num_epochs=num_epochs, is_inception=(model_name==\"inception\"))\n\n# Plot the training curves of validation accuracy vs. number \n#  of training epochs for the transfer learning method and\n#  the model trained from scratch\nohist = []\nshist = []\n\nohist = [h.cpu().numpy() for h in hist]\nshist = [h.cpu().numpy() for h in scratch_hist]\n\nplt.title(\"Validation Accuracy vs. Number of Training Epochs\")\nplt.xlabel(\"Training Epochs\")\nplt.ylabel(\"Validation Accuracy\")\nplt.plot(range(1,num_epochs+1),ohist,label=\"Pretrained\")\nplt.plot(range(1,num_epochs+1),shist,label=\"Scratch\")\nplt.ylim((0,1.))\nplt.xticks(np.arange(1, num_epochs+1, 1.0))\nplt.legend()\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Final Thoughts and Where to Go Next\n-----------------------------------\n\nTry running some of the other models and see how good the accuracy gets.\nAlso, notice that feature extracting takes less time because in the\nbackward pass we do not have to calculate most of the gradients. There\nare many places to go from here. You could:\n\n-  Run this code with a harder dataset and see some more benefits of\n   transfer learning\n-  Using the methods described here, use transfer learning to update a\n   different model, perhaps in a new domain (i.e.\u00a0NLP, audio, etc.)\n-  Once you are happy with a model, you can export it as an ONNX model,\n   or trace it using the hybrid frontend for more speed and optimization\n   opportunities.\n\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.8"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}